package com.xingchi.tornado.core.config;

import com.xingchi.tornado.core.constants.Constants;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.MDC;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.TaskDecorator;
import org.springframework.lang.NonNull;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * 线程池配置
 *
 * @author xingchi
 * @date 2024/5/30 8:55
 */
@Slf4j
@Configuration
public class ExecutePoolConfiguration {

    @Bean
    public ThreadPoolTaskExecutor defaultExecutorService() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        int availableProcessors = Runtime.getRuntime().availableProcessors();
        int corePoolSize = availableProcessors * 2;
        int maxPoolSize = availableProcessors * 4;
        int queueCapacity = availableProcessors * 100;
        executor.setCorePoolSize(corePoolSize);
        executor.setMaxPoolSize(maxPoolSize);
        executor.setQueueCapacity(queueCapacity);
        executor.setKeepAliveSeconds(60);
        executor.setTaskDecorator(new TraceIdTaskDecorator());
        executor.setThreadNamePrefix("default-thread-");
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        executor.initialize();
        log.info("IO denseness computation  thread pool init successfully! corePoolSize:{} maxPoolSize:{} queueCapacity:{}", corePoolSize, maxPoolSize, queueCapacity);
        return executor;
    }

    @Bean
    public ThreadPoolTaskExecutor cpuThreadPool() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        int availableProcessors = Runtime.getRuntime().availableProcessors();
        executor.setCorePoolSize(availableProcessors);
        int maxPoolSize = availableProcessors * 2;
        int queueCapacity = availableProcessors * 100;
        executor.setMaxPoolSize(maxPoolSize);
        executor.setQueueCapacity(queueCapacity);
        executor.setKeepAliveSeconds(30);
        executor.setTaskDecorator(new TraceIdTaskDecorator());
        executor.setThreadNamePrefix("cpu-thread-");
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        executor.initialize();
        log.info("CPU denseness computation thread pool init successfully! corePoolSize:{} maxPoolSize:{} queueCapacity:{}", executor.getCorePoolSize(), executor.getMaxPoolSize(), queueCapacity);
        return executor;
    }

    class TraceIdTaskDecorator implements TaskDecorator {

        @Override
        @NonNull
        public Runnable decorate(@NonNull Runnable runnable) {
            Map<String, String> mdcMappings = Optional.ofNullable(MDC.getCopyOfContextMap()).orElse(new HashMap<>());
            return () -> {
                String traceId = mdcMappings.get(Constants.MDC_TRACE_ID);
                if (StringUtils.isBlank(traceId)) {
                    traceId = getTraceId();
                    log.debug("TraceId cannot be found in MDC! Needs to be regenerated new TraceId:{}", traceId);
                }

                MDC.put(Constants.MDC_TRACE_ID, traceId);
                try {
                    runnable.run();
                } finally {
                    MDC.remove(Constants.MDC_TRACE_ID);
                }
            };
        }
    }

    public String getTraceId() {
        return UUID.randomUUID().toString().replace("-", "");
    }


}
